
# - Step 3 (freq.py): collecting sucess-frequencies for the generated Insights
# 	- For each question + reasoning history, generate a reasoning-situation
# 	- Use the embedding of the reasoning-situation to find K most similar Insights
# 	- Randomly select one Insight as example to generate a insight for the curren the reasoning-situation
# 	- Use Insights generated this way to generate the complete reasoning for the question

from com import create, emb
from exp import data, run
from util import color, py_solver
import os, random
from datetime import datetime

def _get_asnwer(model_id, question, reasoning_steps, verbal=False):
	resp = run.generate(model_id, question, 'com/reasoner_prompts', 'answer',
						question=question.question, 
						reasoning_steps=reasoning_steps, 
						asnwer_format=question.asnwer_format,
						verbal=verbal)
	target, answer = resp.strip().split('\n', 1)
	target_prefix = 'The meaning of the answer is: '
	assert target.startswith(target_prefix)
	target = target[len(target_prefix):].strip()
	return target, answer

def strip_markdown(resp):
	if '```markdown' in resp:
		resp_ = py_solver.extract_block(resp, 'markdown')
		assert resp_ is not None, resp
		resp = resp_
	return resp

def _get_insight(model_id, question, reasoning_steps, verbal, examples=None):
	examples_prefix = 'insight_example' if examples is None else None
	resp = run.generate(model_id, question, 'com/reasoner_prompts', 'insight',
						examples_prefix=examples_prefix,
						question=question.question, 
						examples=examples, 
						reasoning_steps='\n\n'.join(reasoning_steps),
						step=len(reasoning_steps)+1,
						verbal=verbal,
						)
	resp = strip_markdown(resp)
	if '- Reasoning Step ' in resp:
		resp = resp.split('- Reasoning Step ', 1)[0]
	insights = create.parse_insights(resp, '- Insight For Reasoning Step')
	return insights[0]

insight_repr_templ = '''- Insight For Reasoning Step {step}:
  - Situation: {situation}
  - Goal: {goal}'''

def _reasoning_steps_repr(reasoning_steps, last_step=None):
	steps = []
	for i, reasoning_step in enumerate(reasoning_steps):
		insight = reasoning_step['insight']
		reasoning = reasoning_step['reasoning']
		step = [insight_repr_templ.format(step=i+1,
										situation=insight['Situation'],
										goal=insight['Goal'],
										)]
		if (last_step is None) or (i < last_step):
			step.append('')
			step.append(f'- Reasoning Step {i+1}: {reasoning}')
		steps.append('\n'.join(step))
		if (last_step is not None) and (i == last_step):
			break
	return steps

def _examples_repr(examples, insight_examples):
	def _make_example(e, j):
		problem = e['problem']
		reasoning_steps = [{'insight':i, 'reasoning':r} for i, r in zip(e['insights'], e['reasoning_steps'])]
		assert j < len(reasoning_steps)
		reasoning = _reasoning_steps_repr(reasoning_steps, last_step=j)
		reasoning = '\n\n'.join(reasoning)
		return f'**Problem**: {problem}\n\n**Reasoning**:\n\n{reasoning}'
	examples_ = [_make_example(examples[i], j) for i, j in insight_examples]
	return '\n\n'.join([f'## Example {i+1}\n\n{e}' for i,e in enumerate(examples_)])

def _get_reasoning_step(model_id, question, reasoning_steps, verbal, insight):
	step = len(reasoning_steps) + 1
	insight = insight_repr_templ.format(step=step,
										situation=insight['Situation'],
										goal=insight['Goal'],
										)
	resp = run.generate(model_id, question, 'com/reasoner_prompts', 'reasoning',
						examples_prefix='reasoing_example',
						question=question.question, 
						reasoning_steps='\n\n'.join(reasoning_steps),
						step=step,
						insight=insight,
						verbal=verbal,
						)
	resp = strip_markdown(resp)
	reasoning_steps = create.parse_reasoning_steps(resp, '- Reasoning Step')
	reasoning_step = reasoning_steps[-1]
	if '- End of Reasoning:' in reasoning_step:
		reasoning_step, end_of_reasoning = reasoning_step.split('- End of Reasoning:', 1)
	else:
		reasoning_step, end_of_reasoning = reasoning_step, ''
	reasoning_step = reasoning_step.lstrip(' ').rstrip()
	end_of_reasoning = 'True' in end_of_reasoning
	return reasoning_step, end_of_reasoning

def _run_code(resp, verbal):
	resp = resp.rstrip()
	code = py_solver.extract_block(resp, 'python')
	if code is None:
		return resp
	try:
		output = py_solver.execute_code(code).rstrip()
		if verbal:
			color.text(f'\nPython output: {output}', 'blue')
		if output is None:
			output = 'None'
		output = output.strip()
		if output == '':
			output = '<Empty>'
		resp += '\n  - Python code outputs: ' + output
	except Exception as ex:
		exc = str(ex)
		if verbal:
			color.text('\n'+exc, 'red')
		resp += '\n  - Python code raise an exception:\n' + exc.strip('\n').replace('\n\n', '\n')
	return resp

class Reasoner:
	def __init__(self, examples, embs, emb_index, 
				topk=25, num_insight_examplars=1, 
				max_steps=10, enable_coding=True):
		self.examples, self.embs, self.emb_index = examples, embs, emb_index
		self.topk = topk
		self.num_insight_examplars = num_insight_examplars
		self.max_steps = max_steps
		self.enable_coding = enable_coding

	def _retrieve_insight_example(self, situation, pid):
		s_emb = emb.gen_emb(situation)
		insights = emb.retrieve_insights(s_emb, pid, self.examples, self.embs, self.emb_index, self.topk)
		return random.sample(insights, self.num_insight_examplars)

	def _step(self, model_id, question, reasoning_steps, verbal):
		steps_repr = _reasoning_steps_repr(reasoning_steps)
		insight = _get_insight(model_id, question, steps_repr, verbal)
		insight_examples = None
		if self.topk > 0 and self.num_insight_examplars > 0:
			insight_examples = self._retrieve_insight_example(insight['Situation'], question.pid)
			examples_repr = _examples_repr(self.examples, insight_examples)
			insight = _get_insight(model_id, question, steps_repr, verbal, examples_repr)
		reasoning_step, end_of_reasoning = _get_reasoning_step(model_id, question, steps_repr, verbal, insight)
		if self.enable_coding:
			reasoning_step = _run_code(reasoning_step, verbal)
		reasoning_steps.append({
			'insight': insight,
			'reasoning': reasoning_step,
			'examples': insight_examples,
			})
		return end_of_reasoning

	def reason(self, model_id, question, verbal=True):
		reasoning_steps = []
		for _ in range(self.max_steps):
			if self._step(model_id, question, reasoning_steps, verbal):
				break
		target, answer = _get_asnwer(model_id, question, '\n\n'.join(_reasoning_steps_repr(reasoning_steps)), verbal)
		return question.extra_asnwer(answer), target, reasoning_steps

examples_path = '../OUT/created-examples'
freq_path = '../OUT/freq'
out_path = '../OUT/TBYS'

def load_data_1(test_set):
	example_paths = [f'{examples_path}/step1_{test_set}.jsonl']
	emb_path = f'{examples_path}/step2_{test_set}_emb.jsonl'
	examples = emb.load_examples(example_paths)
	embs, emb_index = emb.load_emb(emb_path)
	return examples, embs, emb_index

def load_data_3(test_set, num_examples=None):
	examples = emb.load_jsonl(f'{freq_path}/examples_{test_set}.jsonl')
	if num_examples is not None:
		examples = examples[:num_examples]
	print(f'# examples:', len(examples))
	embs, emb_index = emb.load_emb(f'{freq_path}/emb_{test_set}.jsonl', len(examples))
	return examples, embs, emb_index

def test(test_set, example_test_set='MATH500'):
	num_insight_examplars = 3
	try:
		examples, embs, emb_index = load_data_3(example_test_set, 100)
	except:
		try:
			examples, embs, emb_index = load_data_1(example_test_set)
		except:
			examples = create.seed_examples() 
			embs, emb_index = None, None
			num_insight_examplars = 0
	reasoner = Reasoner(examples, embs, emb_index, num_insight_examplars=num_insight_examplars)
	test_examples = data.get_data('MATH500')
	for e in test_examples[26:27]:
		run.test(e, reasoner.reason)

def exp_1(test_set, k_shot, example_size, num_samples=8, example_test_set='MATH500', enable_coding=True, outfile=None):
	if outfile is None:
		outfile = f'{out_path}/{test_set}_{example_size}-examples_{k_shot}-shot.jsonl'
	print(f'# {outfile} {datetime.now()}')
	# if os.path.isfile(outfile): return
	test_examples = data.get_data(test_set)
	examples, embs, emb_index = load_data_3(example_test_set, example_size)
	reasoner = Reasoner(examples, embs, emb_index, num_insight_examplars=k_shot, 
						enable_coding=enable_coding)
	run.run_exp(test_examples, reasoner.reason, num_samples=num_samples, max_retries=num_samples*3, outfile=outfile)

def exp2(test_set, default_example_size=50):
	# exp_1(test_set, 8, default_example_size)

	# # tbys with different examples library size
	for example_size in ([50] + list(range(100, 1100, 100))):
		exp_1(test_set, 8, example_size, 1)

	# tbys with different examples-shot (TBYS MATH+AIME)
	# for k_shot in range(8):
	# 	exp_1(test_set, k_shot, default_example_size, 1)
		
	# - ablation (TBYS MATH+AIME)
	# 	- Use of library (k_shot=0 Done)
	# 	- use of coding
	outfile = f'{out_path}/{test_set}_50-examples_8-shot-nocode.jsonl'
	exp_1(test_set, 8, default_example_size, 1, enable_coding=False, outfile=outfile)

def main(test_set='MATH500'):
	if not os.path.isdir(out_path):
		os.mkdir(out_path)
	for example_size in [50, 100, 150, 200, 250, 300, 350, 400, 450, 500]:
		exp_1(test_set, 8, example_size, 8)

if __name__ == '__main__':
	exp_1('MATH500', 8, 500, 8)
	# test('MATH500')
	# main('MATH500')
	# main('AIME')
	# exp('MATH500')
	# exp('AIME')
	# exp2('MATH500')
	# exp2('AIME', 500)
